/**
 * @file matmul_leakyrelu_custom_tiling.h
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#ifndef QUANT_MATMUL_V2_TILING_H
#define QUANT_MATMUL_V2_TILING_H

#include "register/tilingdata_base.h"
#include "tiling/tiling_api.h"

namespace optiling {
BEGIN_TILING_DATA_DEF(QuantMatmulV2Params)
TILING_DATA_FIELD_DEF(uint32_t, is_per_tensor);
TILING_DATA_FIELD_DEF(uint32_t, is_pertoken);
TILING_DATA_FIELD_DEF(uint32_t, ub_calc_m);
TILING_DATA_FIELD_DEF(uint32_t, ub_calc_n);
TILING_DATA_FIELD_DEF(uint32_t, need_ub_buffer);
TILING_DATA_FIELD_DEF(uint32_t, real_single_core_m);
TILING_DATA_FIELD_DEF(uint32_t, real_single_core_n);
TILING_DATA_FIELD_DEF(uint32_t, bias_dtype);
TILING_DATA_FIELD_DEF(uint32_t, split_trans_num);
TILING_DATA_FIELD_DEF(uint32_t, rank_size);
TILING_DATA_FIELD_DEF(uint32_t, rank);
TILING_DATA_FIELD_DEF(uint32_t, p_value);
END_TILING_DATA_DEF;

REGISTER_TILING_DATA_CLASS(QuantMatmulV2ParamsOp, QuantMatmulV2Params);

BEGIN_TILING_DATA_DEF(L2cacheParams)
TILING_DATA_FIELD_DEF(uint32_t, m_tile_cnt_l2);
TILING_DATA_FIELD_DEF(uint32_t, n_tile_cnt_l2);
TILING_DATA_FIELD_DEF(uint32_t, m_tile_block);
TILING_DATA_FIELD_DEF(uint32_t, n_tile_block);
TILING_DATA_FIELD_DEF(uint32_t, cal_order);
TILING_DATA_FIELD_DEF(uint32_t, is_basic_tiling);
END_TILING_DATA_DEF;

REGISTER_TILING_DATA_CLASS(L2cacheParamsOp, L2cacheParams);

BEGIN_TILING_DATA_DEF(QuantMatmulV2TilingData)
TILING_DATA_FIELD_DEF_STRUCT(QuantMatmulV2Params, params);
TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, matmul_tiling);
TILING_DATA_FIELD_DEF_STRUCT(L2cacheParams, tile_l2cache_tiling);
END_TILING_DATA_DEF;

REGISTER_TILING_DATA_CLASS(QuantMatmulV2, QuantMatmulV2TilingData);
REGISTER_TILING_DATA_CLASS(QuantMatmulV2TilingDataOp, QuantMatmulV2TilingData);

constexpr uint64_t BASIC_ALIGN_16 = 16;
constexpr uint64_t UINT32_MASK = 0x00000000ffffffff;

struct QuantMatmulV2Info{
public:
    uint64_t getMatmulApiMSize();
    uint64_t getTotalMatmulApiMSize(uint64_t base_m);
    uint64_t getTotalBaseMCnt(uint64_t base_m);
    bool trans_a = false;
    bool trans_b = false;
    bool has_bias = false;
    bool is_basic_tiling = false;
    bool is_decode_opt = false;
    uint64_t m_size = 0UL;
    uint64_t m_size_per_npu = 0UL;
    uint64_t n_size = 0UL;
    uint64_t k_size = 0UL;
    uint64_t split_trans_num = 0UL;
    ge::DataType a_dtype = ge::DT_INT8;
    ge::DataType b_dtype = ge::DT_INT8;
    ge::DataType c_dtype = ge::DT_BF16;
    ge::DataType bias_dtype = ge::DT_INT32;
    ge::DataType scale_dtype = ge::DT_BF16;
    bool is_per_tensor = false;
    bool is_pertoken = false;
    int64_t out_dtype = 0L;
    uint32_t lib_api_work_space_size = 16 * 1024 * 1024;
    uint64_t bf16_extre_work_space_size = 0UL;
    const char *op_name = nullptr;
    ge::Format a_format = ge::FORMAT_ND;
    ge::Format b_format = ge::FORMAT_FRACTAL_NZ;
    ge::Format c_format = ge::FORMAT_ND;
};

struct BasicTiling{
    uint64_t used_core_num = 1;
    uint64_t single_core_m = 1;
    uint64_t single_core_n = 1;
    uint64_t single_core_k = 1;
    uint64_t base_m = 1;
    uint64_t base_n = 1;
    uint64_t base_k = 1;
    uint64_t step_ka = 1;
    uint64_t step_kb = 1;
    uint64_t depth_a1 = 1;
    uint64_t depth_b1 = 1;
    uint64_t step_m = 1;
    uint64_t step_n = 1;
    uint64_t iterate_order = 0;
    uint64_t db_l0c = 1;
    uint64_t is_mclash = 0;
    uint64_t is_nclash = 0;
    uint64_t m_tile_cnt_l2 = 0;
    uint64_t n_tile_cnt_l2 = 0;
    uint64_t m_tile_block = 0;
    uint64_t n_tile_block = 0;
    uint64_t cal_order = 0;
};

template<typename T>
T get_shape_with_data_type(T size, ge::DataType dtype) {
    if (dtype == ge::DT_INT4) {
        return size + size;
    } else {
        return size / static_cast<T>(ge::GetSizeByDataType(dtype));
    }
}

template<typename T>
T get_size_with_data_type(T shape, ge::DataType dtype) {
    if (dtype == ge::DT_INT4) {
        return (shape + 1) >> 1;
    } else {
        return shape * static_cast<T>(ge::GetSizeByDataType(dtype));
    }
}

static constexpr uint64_t MB_SIZE = 1024 * 1024;
static constexpr uint64_t KB_SIZE = 1024;
static constexpr uint32_t ROW_FIRST = 1;
static constexpr uint32_t COL_FIRST = 2;
static constexpr uint64_t NUM_DB = 2;
static constexpr uint64_t NUM_DB_MM = 4;

}  // namespace optiling

#endif // QUANT_MATMUL_V2_TILING_H